'''这个代码文件已废弃'''

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

"""这个文件关于不同方法的管理"""

class ConversationManager:
    def __init__(self,tokenizer, model, enti_num=5):
        # def __init__(self, model_name="meta-llama/Llama-2-7b-chat-hf", enti_num=5):
        # self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        # self.model = AutoModelForCausalLM.from_pretrained(
        #     model_name, 
        #     device_map="auto", 
        #     torch_dtype=torch.float16
        # )
        self.tokenizer = tokenizer  # 直接使用传入的 tokenizer
        self.model = model          # 直接使用传入的 model       
        self.enti_num = enti_num
        self.conversation_history = []
    
    def generate_response(self, prompt, max_new_tokens=2000):
        """生成回复并保存对话历史"""
        # 准备输入（移除Llama不支持的token_type_ids）
        inputs = self.tokenizer(prompt, return_tensors="pt", return_token_type_ids=False)
        # return_tensors="pt" 指定返回 PyTorch 张量（"pt" 代表 PyTorch）。若不指定，默认返回 Python 列表。
        inputs = inputs.to(self.model.device)
        
        # 生成配置
        generation_args = {
            **inputs,
            "max_new_tokens": max_new_tokens,
            "output_scores": True,
            "return_dict_in_generate": True,
            "do_sample": True,
            "temperature": 0.7,
            "top_p": 0.9
        }
        
        # 生成回复
        outputs = self.model.generate(**generation_args)
        
        # 解码回复
        generated_token_ids = outputs.sequences[0, inputs.input_ids.shape[1]:]
        response = self.tokenizer.decode(generated_token_ids, skip_special_tokens=True)
        
        # 保存对话
        self.conversation_history.append({
            "prompt": prompt,
            "response": response,
            "scores": outputs.scores
        })
        
        return response

    def entity_extraction_method(self, biography):
        """实体提取法（4轮对话）"""
        # 第一轮
        prompt1 = f"从生成的传记中提取 {self.enti_num} 个最重要的实体，并在该传记中标注出来。\n传记：{biography}"
        print("prompt1",prompt1)
        # response1 = self.generate_response(prompt1)
        response1 = manager.generate_full_response(prompt1)

        print("response1",response1)
        '''
        # 第二轮
        prompt2 = f"依据这 {self.enti_num} 个实体，将上述传记包含的事实拆解为 {self.enti_num} 个部分，每个部分均以 {{<实体>,<对应事实描述>}} 的形式呈现。"
        # print("prompt2",prompt2)
        response2 = self.generate_response(prompt2)
        # print("response2",response2)
        
        # 第三轮
        prompt3 = f"重新核查所生成的传记，确认围绕这 {self.enti_num} 个实体本身及针对这些实体展开的描述是否正确，若存在错误，需找出原因。"
        # print("prompt3",prompt3)
        response3 = self.generate_response(prompt3)
        # print("response3",response3)

        # 第四轮
        prompt4 = "仔细检查每个实体及其对应的事实描述的正确性，涵盖时间、地点及其他实体间的关系等细节，纠正错误后，重新生成 {<实体>，<对应事实描述>} 的要点，并最终重新撰写传记。"
        # print("prompt4",prompt4)
        response4 = self.generate_response(prompt4)
        # print("response4",response4)
'''
        return self.conversation_history[-4:]  # 返回最后四轮对话

    def one_reflection_method(self, biography):
        """一次反思法（1轮对话）"""
        prompt = f"仔细检查每个实体及其对应的事实描述的正确性，涵盖时间、地点及其他实体间的关系等细节，纠正错误后，重新生成 {{<实体>，<对应事实描述>}} 的要点，并最终重新撰写传记。\n传记：{biography}"
        response = self.generate_response(prompt)
        
        return self.conversation_history[-1:]  # 返回最后一轮对话

    def atomic_method(self, biography):
        """原子方法（2轮对话）"""
        # 第一轮
        prompt1 = f"将上述内容拆解为一个个独立且不可分割的原子事实，一一列出。\n传记：{biography}"
        response1 = self.generate_response(prompt1)
        
        # 第二轮
        prompt2 = "检查每一个事实的正确性，修改纠正后，生成修改后的传记。再次仔细检查每个实体及其对应的事实描述的正确性，包括时间、地点和其他实体间的关系等细节，修正答案。"
        response2 = self.generate_response(prompt2)
        
        return self.conversation_history[-2:]  # 返回最后两轮对话




"""这里是主函数"""
# 使用示例
if __name__ == "__main__":
    # 初始化对话管理器
    manager = ConversationManager(model_name="meta-llama/Llama-2-7b-chat-hf", enti_num=5)
    
    # 示例传记文本
    sample_biography = "艾萨克·牛顿（1643年1月4日—1727年3月31日）是英国著名的物理学家、数学家、天文学家、自然哲学家和炼金术士。"
    
    # 运行实体提取法
    print("===== 运行实体提取法 =====")
    entity_results = manager.entity_extraction_method(sample_biography)
    
    # 运行一次反思法
    print("\n===== 运行一次反思法 =====")
    reflection_results = manager.one_reflection_method(sample_biography)
    
    # 运行原子方法
    print("\n===== 运行原子方法 =====")
    atomic_results = manager.atomic_method(sample_biography)
    
    # 保存完整对话历史
    manager.save_conversation()
    
    # 打印最后一次对话
    print("\n最后一次对话内容：")
    last_entry = manager.conversation_history[-1]
    print(f"用户输入：\n{last_entry['prompt']}\n")
    print(f"模型回复：\n{last_entry['response']}\n")